# from torch import nn
#
#
# class MyModule(nn.Module):
#     def __init__(self, *args, **kwargs) -> None:
#         super().__init__(*args, **kwargs)
#     def forward(self,input):
#         output = input + 1
#         return output
# test_module = MyModule()
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
import torch.nn.functional as F

class MyModule(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
    def forward(self,input):
        output = input + 1
        return output

class Mydata(Dataset):
    def __init__(self, root_dir, label_dir, transform=None):  # 新增 transform 参数
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.image_path = os.listdir(self.path)
        self.transform = transform  # 保存 transform 用于后续处理

    def __getitem__(self, idx):
        image_name = self.image_path[idx]
        image_item_path = os.path.join(self.path, image_name)
        image = Image.open(image_item_path)
        label = self.label_dir

        # 关键：在数据集内部对单个图像应用 transform（转为 Tensor）
        if self.transform:
            image = self.transform(image)

        return image, label  # 此时 image 已经是 Tensor

    def __len__(self):
        return len(self.image_path)


# 定义 transform（Resize + 转 Tensor）
dataset_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

train_root_dir = "train"
ants_label_dir = "ants"
# 创建数据集时传入 transform，确保返回的是 Tensor
ants_dataset = Mydata(train_root_dir, ants_label_dir, transform=dataset_transform)
# DataLoader 加载的 batch 中，图像已经是 Tensor
test_loader = DataLoader(dataset=ants_dataset, batch_size=4, shuffle=True, num_workers=0, drop_last=False)

# 错误示范：
# kernel = torch.tensor([[1,2,1],
#                        [0,1,0],
#                        [2,1,0]])
# kernel = torch.reshape(kernel,(1,3,3,3))

# 调整卷积核：in_channels=3（与输入图像通道数匹配）
# 方法：将1通道卷积核复制3次，适应3通道输入（也可自定义3个不同的核）
kernel = torch.tensor([[1,2,1],
                       [2,1,0],
                       [1,1,2]], dtype=torch.float32)
# 重塑为[1, 3, 3, 3]：out_channels=1，in_channels=3，核大小3x3
kernel = kernel.unsqueeze(0).unsqueeze(0)  # 先转为[1,1,3,3]
kernel = kernel.repeat(3, 3, 1, 1)  # 复制3次，变为[3,3,3,3]（in_channels=3,out_channels=3）

writer = SummaryWriter("loader_logs")
step = 0
for data in test_loader:
    ants_imgs, ants_labels = data  # ants_imgs 是 shape 为 [4, 3, 256, 256] 的 Tensor（批量图像）
    ants_imgs_conv2d = F.conv2d(ants_imgs,kernel,stride=1)
    writer.add_images("test",ants_imgs_conv2d,step)
    step += 1
writer.close()
